import json
import time
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Union, Optional, List

from torchvision import transforms

# 导入自定义模块
from yms_zsl.tools.tools import setup_save_dirs, setup_logger, get_device
from yms_zsl.train import main as main_train
from yms_zsl.train.main import TrainConfig
from yms_zsl.train.predict import PredictConfig, predict
from yms_zsl.train.semantics import SemanticsConfig, semantics
from yms_zsl.train.train_fcnn import TrainFCnnConfig, train_fcnn


def save_experiment_config(config: dataclass, save_dir: Union[str, Path],
                           filename: str = "experiment_config.json") -> None:
    """保存实验配置到JSON文件，方便复现（修复路径拼接错误）"""
    # 核心修改：将 save_dir 统一转为 Path 对象（支持 / 拼接）
    save_dir = Path(save_dir)
    config_dict = asdict(config)

    # 转换Path对象为字符串（JSON不支持Path类型）
    for k, v in config_dict.items():
        if isinstance(v, Path):
            config_dict[k] = str(v.absolute())
        elif isinstance(v, list) and len(v) > 0 and isinstance(v[0], Path):
            config_dict[k] = [str(p.absolute()) for p in v]

    # 现在 save_dir 是 Path 对象，可正常用 / 拼接
    save_path = save_dir / filename
    with open(save_path, "w", encoding="utf-8") as f:
        json.dump(config_dict, f, ensure_ascii=False, indent=4)
    # 打印完整路径，方便查看
    print(f"📝 实验配置已保存至：{save_path.absolute()}")


def run(configs: "Configs") -> None:
    # -------------------------- 0. 初始化：保存目录+日志+设备 --------------------------
    try:
        # 1. 初始化保存目录（不管 setup_save_dirs 返回字符串还是Path，后续都会统一处理）
        save_dir = setup_save_dirs(configs.save_dir, prefix="exp")
        logger = setup_logger(save_dir)

        # 2. 保存实验配置（此时 save_dir 可能是字符串，函数内部会转为Path）
        save_experiment_config(configs, save_dir)

        # 3. 设备初始化
        device = get_device()
        logger.info("=" * 60)
        logger.info("🎉 零样本学习完整训练流程启动")
        logger.info(f"📅 启动时间：{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}")
        logger.info(f"💻 计算设备：{device.type}（ID: {device.index if device.index is not None else '无'}）")
        logger.info(f"📁 结果保存目录：{Path(save_dir).absolute()}")  # 转为Path打印完整路径
        logger.info(f"⚠️  提示：请确认 data_dir={configs.data_dir} 和 train_class={configs.train_class} 路径正确")
        logger.info("=" * 60)

        # 4. 统一图像变换
        transform = transforms.Compose([
            transforms.Resize((32, 32)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        logger.info("\n🖼️  图像预处理管道：Resize(32,32) → ToTensor → Normalize")

        # -------------------------- 1. 阶段1/4：特征解耦训练（CNN + DECAE） --------------------------
        logger.info("\n" + "=" * 50)
        logger.info("📌 阶段1/4：特征提取网络（CNN）+ 卷积自编码器（DECAE）训练")
        logger.info("=" * 50)

        main_config = TrainConfig(
            device=device,
            transform=transform,
            data_dir=str(configs.data_dir),
            save_dir=str(save_dir),  # 若需要字符串，转为str（避免Path不兼容）
            train_class=str(configs.train_class),
            batch_size=configs.batch_size,
            num_workers=configs.num_workers,
            cnn_lr=configs.cnn_lr,
            decae_lr=configs.decae_lr,
            epochs=configs.epochs,
        )
        logger.info(f"阶段配置：CNN_LR={configs.cnn_lr}, DECAE_LR={configs.decae_lr}, Epochs={configs.epochs}")

        # 训练执行+耗时统计
        start_time = time.perf_counter()
        main_train.main(main_config)
        end_time = time.perf_counter()
        stage1_time = end_time - start_time
        logger.info(f"✅ 阶段1完成！耗时：{stage1_time:.2f}秒（{stage1_time / 60:.2f}分钟）")

        # -------------------------- 2. 阶段2/4：属性提取（语义构建） --------------------------
        logger.info("\n" + "=" * 50)
        logger.info("📌 阶段2/4：属性提取与语义构建")
        logger.info("=" * 50)

        semantics_config = SemanticsConfig(
            device=device,
            transform=transform,
            data_dir=str(configs.data_dir),
            save_dir=str(save_dir),
            train_class=str(configs.train_class),
            ignore_factors=configs.ignore_factors,
            factor_index_map_path=configs.factor_index_map_path,
        )
        logger.info(f"忽略因子：{configs.ignore_factors if configs.ignore_factors else '无'}")

        start_time = time.perf_counter()
        semantics(semantics_config)
        end_time = time.perf_counter()
        stage2_time = end_time - start_time
        logger.info(f"✅ 阶段2完成！耗时：{stage2_time:.2f}秒（{stage2_time / 60:.2f}分钟）")

        # -------------------------- 3. 阶段3/4：FCnn嵌入模型训练 --------------------------
        logger.info("\n" + "=" * 50)
        logger.info("📌 阶段3/4：FCnn嵌入模型训练（特征-语义映射）")
        logger.info("=" * 50)

        fc_config = TrainFCnnConfig(
            device=device,
            transform=transform,
            data_dir=str(configs.data_dir),
            save_dir=str(save_dir),
            train_class=str(configs.train_class),
            batch_size=configs.batch_size,
            epochs=configs.epochs,
            num_workers=configs.num_workers,
            fcnn_channels=configs.fcnn_channels,
            factor_index_map_path=configs.factor_index_map_path,
        )
        logger.info(f"FCnn配置：输入通道={fc_config.fcnn_channels}，BatchSize={fc_config.batch_size}")

        start_time = time.perf_counter()
        train_fcnn(fc_config)
        end_time = time.perf_counter()
        stage3_time = end_time - start_time
        logger.info(f"✅ 阶段3完成！耗时：{stage3_time:.2f}秒（{stage3_time / 60:.2f}分钟）")

        # -------------------------- 4. 阶段4/4：测试集预测（无返回值适配） --------------------------
        logger.info("\n" + "=" * 50)
        logger.info("📌 阶段4/4：测试集零样本预测")
        logger.info("=" * 50)

        # 构建预测配置（路径统一转为str，避免兼容问题）
        predict_config = PredictConfig(
            device=device,
            transform=transform,
            data_dir=str(configs.data_dir),
            save_dir=str(save_dir),
            batch_size=configs.batch_size,
            fcnn_channels=configs.fcnn_channels,
            factor_index_map_path=configs.factor_index_map_path,
            train_class=configs.predict_class,
            ignore_factors=configs.ignore_factors,
        )
        # 执行预测（无返回值，直接调用）
        start_time = time.perf_counter()
        predict(predict_config)
        end_time = time.perf_counter()
        stage4_time = end_time - start_time
        logger.info(f"✅ 阶段4预测完成！耗时：{stage4_time:.2f}秒（{stage4_time / 60:.2f}分钟）")

        # -------------------------- 5. 整体流程总结 --------------------------
        logger.info("\n" + "=" * 60)
        logger.info("📊 完整流程总结")
        logger.info("=" * 60)

        total_time = stage1_time + stage2_time + stage3_time + stage4_time
        logger.info("⏱️  各阶段耗时明细：")
        logger.info(f"   1. 特征解耦训练：{stage1_time:.2f}秒")
        logger.info(f"   2. 属性提取：{stage2_time:.2f}秒")
        logger.info(f"   3. FCnn训练：{stage3_time:.2f}秒")
        logger.info(f"   4. 测试集预测：{stage4_time:.2f}秒")
        logger.info(f"\n📅 整体总耗时：{total_time:.2f}秒（{total_time / 60:.2f}分钟）")
        logger.info(f"📁 所有结果（含预测输出）保存于：{Path(save_dir).absolute()}")
        logger.info("✅ 零样本学习完整流程全部完成！")
        logger.info("=" * 60)

    except Exception as e:
        # 全局异常捕获，打印错误信息
        logger.error("\n❌ 流程执行失败：", exc_info=True)
        logger.error(f"错误详情：{str(e)}")
        raise


@dataclass
class Configs:
    """零样本学习完整流程配置类（参数集中管理）"""
    # 路径配置（需手动确认正确性）
    data_dir: Union[str, Path] = r'D:\Code\2-ZSL\0-data\CWRU\dataset'
    save_dir: Union[str, Path] = r'D:\Code\2-ZSL\1-output\ZSL_实验结果'
    train_class: Union[str, Path] = r'D:\Code\2-ZSL\0-data\CWRU\dataset\seen_classes.txt'
    predict_class: Union[str, Path] = None
    factor_index_map_path: Optional[Union[str, Path]] = None

    # 训练超参数
    epochs: int = 100
    batch_size: int = 40
    cnn_lr: float = 1e-3
    decae_lr: float = 1e-3
    weight_decay: float = 1e-5
    patience: int = 10

    # 数据加载与模型配置
    num_workers: int = 0
    ignore_factors: Optional[List[str]] = None
    fcnn_channels: int = 518


if __name__ == '__main__':
    # 初始化实验配置（根据实际环境修改）
    opts = Configs(
        data_dir=r'D:\Code\2-ZSL\0-data\data\data',
        save_dir=r'D:\Code\2-ZSL\1-output\论文实验结果\对比方法\CWRU\HSAZLM\H06',
        train_class=r'D:\Code\2-ZSL\0-data\data\data/seen_classes.txt',
        factor_index_map_path=r'D:\Code\2-ZSL\0-data\data\data\factor_index_map.txt',
        epochs=1000,
        batch_size=256,
        ignore_factors=["Operating Condition"],
        predict_class=r'D:\Code\2-ZSL\0-data\data\data\unseen_classes.txt'
    )
    run(opts)